import os
from fastapi import FastAPI,UploadFile
import shutil
import uuid
import datetime
import cv2
import base64
import numpy as np
from wechat_ocr.ocr_manager import OcrManager, OCR_MAX_TASK_ID


#------------------------------------------------------------------------------------Api部分
#微信ocr配置
wechat_ocr_dir = "C:\\Users\\ASUS\\AppData\\Roaming\\Tencent\\WXWork\\WeChatOCR\\1.0.1.28\\WeChatOCR\\WeChatOCR.exe"
wechat_dir = "C:\\Program Files (x86)\\Tencent\\WeChat\\[3.9.11.19]" 
ocr_manager = OcrManager(wechat_dir)
# 设置WeChatOcr目录
ocr_manager.SetExePath(wechat_ocr_dir)
# 设置微信所在路径
ocr_manager.SetUsrLibDir(wechat_dir)
# 启动ocr服务
ocr_manager.StartWeChatOCR()

#COR识别
def ocr(img_path:str):
    ocr_manager.DoOCRTask(img_path)
    while ocr_manager.m_task_id.qsize() != OCR_MAX_TASK_ID:
        pass

#识别成功的回调
def ocr_result_callback(img_path:str, results:dict):    
    fileName = os.path.basename(img_path)
    text = ""
    for item in results["ocrResult"]:  
        text += item['text'] + '\n'
    OcrResultDict[fileName] = text[:-1]#这里去掉最后的换行符号
# 设置ocr识别结果的回调函数
ocr_manager.SetOcrResultCallback(ocr_result_callback)


#------------------------------------------------------------------------------------公共部分
#保存OCR识别结果
OcrResultDict = {}

#统一返回结果
class ResultModel:
    success : bool = False # 是否成功
    message: str = None # 响应信息
    data: str = None # 数据

    def fail(self, message:str):
        self.success = False
        self.message = message
        return self

    def success(self, message:str,data = None):
        self.success = True
        self.message = message
        self.data = data
        return self

#base编码部分
class ImgBase64:
    def base64_to_ndarray(self,b64_data: str):
        """base64转numpy数组
        Args: b64_data (str): base64数据
        Returns: _type_: _description_
        """
        image_bytes = base64.b64decode(b64_data)
        image_np = np.frombuffer(image_bytes, dtype=np.uint8)
        image_np2 = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
        return image_np2


    def bytes_to_ndarray(self,img_bytes: str):
        """字节转numpy数组
        Args:  img_bytes (str): 图片字节
        Returns:_type_: _description_
        """
        image_array = np.frombuffer(img_bytes, dtype=np.uint8)
        image_np2 = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
        return image_np2


#------------------------------------------------------------------------------------功能处理部分
#OCR封装处理器(图片)
def ocrHandler(imgPath:str):
    result = ResultModel()
    try:
        key = getFileName(imgPath)
        ocr(imgPath)   
        result.success("识别成功",OcrResultDict[key])
        del OcrResultDict[key]  
        return result
    except Exception as e:
        return result.fail(e.args[0])
    
#OCR封装处理器(上传文件)
def ocrHandlerFile(file: UploadFile):
    if file.filename.endswith((".jpg", ".png")):  # 只处理常见格式图片
        #图片保存路径
        savePath = getSaveImgPath()
        
        #保存图片
        saveFileName = f"{uuid.uuid4()}{os.path.splitext(file.filename)[-1]}"
        with open(os.path.join(savePath, saveFileName), "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        file.close()

        #Ocr识别
        return ocrHandler(savePath + "/" + saveFileName)
    else:
        return ResultModel().fail("识别上传文件失败,只支持[jpg]、[png]上传文件[" + file.filename + "]")
    
#OCR封装处理器(Base64)
def ocrHandlerBase64(base64_str: str):
    try:
        if(base64_str.startswith('data:image') == False):
            return ResultModel().fail("请检查头部代码[data:image/xxx;base64,]为必须")
        imgType = ""
        if(base64_str.startswith('data:image/jpeg;base64,')):
            imgType = ".jpg"
        elif(base64_str.startswith('data:image/png;base64,')):
            imgType = ".png"
        elif(base64_str.startswith('data:image/bmp;base64,')):
            imgType = ".bmp"
        else:
            return ResultModel().fail("不支持的文件类型")

        #保存图片
        img = ImgBase64().base64_to_ndarray(base64_str.split(',')[1])
        savePath = getSaveImgPath()
            
        #保存图片
        saveFileName = f"{uuid.uuid4()}{imgType}"
        cv2.imwrite(savePath + "/" + saveFileName, img)
        return ocrHandler(savePath + "/" + saveFileName)
    except Exception as e:    
        return ResultModel().fail(e.args[0])

#根据文件路径返回文件名
def getFileName(img_path:str):
    fullPath = ""
    if(img_path.startswith("file/20")):#相对路径
        fullPath = os.path.join(os.getcwd(), img_path)
    else:#绝对路径
        fullPath = img_path

    if(os.path.exists(fullPath)):
        return os.path.basename(fullPath)
    else:
        raise ValueError("文件不存在,请检查[" + fullPath + "]")

#处理图片的保存路径
def getSaveImgPath():
    #图片保存路径
    savePath = f'file/{datetime.datetime.now().strftime('%Y-%m-%d')}'
    if not os.path.exists(savePath):
        os.mkdir(savePath)
    return savePath

#------------------------------------------------------------------------------------Api部分
#api接口配置
api = FastAPI(title="OCR API", description="基于 微信的Ocr识别图片和 和 FastAPI 的接口")
count = 0


@api.get('/ocrByLocal', summary="识别本地图片")
def ocrByLocal(localImagePath: str):
    return ocrHandler(localImagePath)


@api.post('/ocrByUpload',summary="识别上传图片")
def ocrByUpload(file: UploadFile):
    return ocrHandlerFile(file)

@api.post('/ocrByBase64',summary="识别Base64编码图片")
def ocrByBase64(base64_str: str):
    return ocrHandlerBase64(base64_str)



# 直接在doc环境下运行启动
# uvicorn ocr:api --host 0.0.0.0 --port 10086
# 在浏览器中查看结果
# http://localhost:10086/docs